import pygame
import random
import math

# 初始化Pygame
pygame.init()

# 屏幕尺寸
WIDTH, HEIGHT = 800, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("烟花特效")

# 颜色定义
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
COLORS = [
    (255, 0, 0),    # 红色
    (0, 255, 0),    # 绿色
    (0, 0, 255),    # 蓝色
    (255, 255, 0),  # 黄色
    (255, 0, 255),  # 紫色
    (0, 255, 255),  # 青色
    (255, 165, 0),  # 橙色
    (255, 192, 203), # 粉色
]

# 粒子类
class Particle:
    def __init__(self, x, y, color):
        self.x = x
        self.y = y
        self.color = color
        self.radius = random.randint(2, 4)
        self.velocity_x = random.uniform(-5, 5)
        self.velocity_y = random.uniform(-5, 5)
        self.gravity = 0.1
        self.life = 100
        self.decay = random.uniform(0.5, 1.5)
        
    def update(self):
        self.x += self.velocity_x
        self.y += self.velocity_y
        self.velocity_y += self.gravity
        self.life -= self.decay
        
    def draw(self, surface):
        if self.life > 0:
            alpha = min(255, int(255 * (self.life / 100)))
            color_with_alpha = (*self.color, alpha)
            pygame.draw.circle(surface, self.color, (int(self.x), int(self.y)), self.radius)
            
    def is_dead(self):
        return self.life <= 0 or self.y > HEIGHT

# 爱心粒子类
class HeartParticle:
    def __init__(self, x, y, color):
        self.x = x
        self.y = y
        self.color = color
        self.scale = random.uniform(0.5, 1.5)
        self.velocity_x = random.uniform(-3, 3)
        self.velocity_y = random.uniform(-3, 3)
        self.gravity = 0.05
        self.life = 100
        self.decay = random.uniform(0.5, 1.0)
        self.angle = random.uniform(0, 2 * math.pi)
        self.angular_velocity = random.uniform(-0.1, 0.1)
        
    def update(self):
        self.x += self.velocity_x
        self.y += self.velocity_y
        self.velocity_y += self.gravity
        self.life -= self.decay
        self.angle += self.angular_velocity
        
    def draw(self, surface):
        if self.life > 0:
            # 绘制爱心形状
            points = []
            for i in range(0, 360, 10):
                rad = math.radians(i)
                # 爱心方程
                px = 16 * (math.sin(rad) ** 3)
                py = -(13 * math.cos(rad) - 5 * math.cos(2*rad) - 2 * math.cos(3*rad) - math.cos(4*rad))
                # 缩放和旋转
                rotated_x = px * math.cos(self.angle) - py * math.sin(self.angle)
                rotated_y = px * math.sin(self.angle) + py * math.cos(self.angle)
                points.append((self.x + rotated_x * self.scale, self.y + rotated_y * self.scale))
            
            if len(points) > 2:
                pygame.draw.polygon(surface, self.color, points)
            
    def is_dead(self):
        return self.life <= 0 or self.y > HEIGHT

# 烟花类
class Firework:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.particles = []
        self.exploded = False
        self.color = random.choice(COLORS)
        self.particle_count = random.randint(50, 100)
        self.show_heart = random.choice([True, False])  # 50%概率显示爱心
        self.show_text = self.show_heart  # 如果显示爱心，也显示文字
        self.text_timer = 0
        
    def explode(self):
        # 决定使用普通粒子还是爱心粒子
        if self.show_heart:
            for _ in range(self.particle_count):
                self.particles.append(HeartParticle(self.x, self.y, self.color))
        else:
            for _ in range(self.particle_count):
                self.particles.append(Particle(self.x, self.y, self.color))
        self.exploded = True
        
    def update(self):
        if not self.exploded:
            self.explode()
            
        for particle in self.particles[:]:
            particle.update()
            if particle.is_dead():
                self.particles.remove(particle)
                
        # 更新文字计时器
        if self.show_text:
            self.text_timer += 1
                
    def draw(self, surface):
        for particle in self.particles:
            particle.draw(surface)
            
    def draw_text(self, surface, font):
        if self.show_text and self.text_timer < 120:  # 显示2秒(60fps*2)
            text = "卢诗琦"
            text_surface = font.render(text, True, WHITE)
            text_rect = text_surface.get_rect(center=(self.x, self.y - 30))
            surface.blit(text_surface, text_rect)
            
    def is_finished(self):
        return self.exploded and len(self.particles) == 0

# 主函数
def main():
    clock = pygame.time.Clock()
    fireworks = []
    running = True
    
    # 创建字体对象
    font = pygame.font.Font(None, 36)
    
    # 创建一个带透明度的surface用于绘制粒子
    particle_surface = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)
    
    while running:
        # 处理事件
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.MOUSEBUTTONDOWN:
                # 鼠标点击时创建烟花
                x, y = pygame.mouse.get_pos()
                fireworks.append(Firework(x, y))
                
        # 更新烟花
        for firework in fireworks[:]:
            firework.update()
            if firework.is_finished():
                fireworks.remove(firework)
                
        # 绘制
        screen.fill(BLACK)
        
        # 在透明surface上绘制粒子
        particle_surface.fill((0, 0, 0, 0))
        for firework in fireworks:
            firework.draw(particle_surface)
            
        # 将粒子surface绘制到屏幕上
        screen.blit(particle_surface, (0, 0))
        
        # 绘制文字
        for firework in fireworks:
            firework.draw_text(screen, font)
        
        pygame.display.flip()
        clock.tick(60)
        
    pygame.quit()

if __name__ == "__main__":
    main()
